import json
import os

from tqdm import tqdm


def prepare_label_studio_data(
    audio_dir, transcription_dir, output_file, audio_url_prefix=""
):
    """
    Prepare data for Label Studio by matching audio files with their transcriptions.

    Parameters:
    audio_dir (str): Directory containing audio files
    transcription_dir (str): Directory containing transcription files
    output_file (str): Path to output JSON file for Label Studio
    audio_url_prefix (str): Prefix for audio URLs (e.g., "/data/local-files/?d=")
    """
    tasks = []
    audio_extensions = [".wav", ".mp3", ".flac", ".ogg", ".m4a"]
    # TODO hack here
    audio_dir_folder_name = os.path.basename(os.path.normpath(audio_dir))

    # Walk through the audio directory
    for root, _, files in os.walk(audio_dir):
        for file in tqdm(files, desc="Processing files"):
            if any(file.lower().endswith(ext) for ext in audio_extensions):
                # Get the relative path from the audio directory
                rel_path = os.path.relpath(os.path.join(root, file), audio_dir)

                # Construct the path to the corresponding transcription file
                base_name = os.path.splitext(file)[0]
                transcription_path = os.path.join(
                    transcription_dir, os.path.dirname(rel_path), f"{base_name}.csv"
                )

                # Check if transcription exists
                if os.path.exists(transcription_path):
                    # Read the CSV file
                    with open(transcription_path, "r", encoding="utf-8") as f:
                        lines = f.readlines()
                        if len(lines) > 1:  # Ensure there's at least a header and data line
                            # Parse CSV data (skip header)
                            language, language_prob, text = lines[1].strip().split(',', 2)
                            # Remove quotes from language and text if present
                            language = language.strip('"')
                            text = text.strip('"')

                    # Construct the audio URL with prefix if provided
                    audio_url = audio_url_prefix + os.path.join(
                        audio_dir_folder_name, rel_path
                    )

                    # Create a task for Label Studio with additional language info
                    task = {
                        "data": {
                            "audio": audio_url,
                            "text": text,
                            "language": language,
                            "language_probability": float(language_prob)
                        }
                    }
                    tasks.append(task)

    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Write tasks to output file
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(tasks, f, ensure_ascii=False, indent=2)

    print(f"Prepared {len(tasks)} tasks for Label Studio in {output_file}")
    return len(tasks)


def export_label_studio_config():
    """
    Return a recommended Label Studio configuration for audio transcription tasks
    """
    config = """
<View>
  <Header value="Transcribe the audio" />
  <Audio name="audio" value="$audio" />
  <Header value="Language Information:"/>
  <View style="display: flex; align-items: center; gap: 10px;">
    <Text name="detected_language" value="Detected language:" />
    <Text name="language_value" value="$language" />
    <Text name="confidence_value" value="(confidence: $language_probability)" />
  </View>
  <TextArea name="language" toName="audio"
            value="$language" rows="1" editable="true"
            placeholder="Edit language code if needed"/>
  <Header value="Current transcription:"/>
  <Text name="current_transcription" value="$text" />
  <Header value="Enter the corrected transcription:"/>
  <TextArea name="transcription" toName="audio" 
            value="$text" rows="5" editable="true" 
            maxSubmissions="1" showSubmitButton="true" />
</View>
"""
    return config.strip()